using System.Collections;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;

public class MTCLTreeSearch
{
    public int simulateCount;
    private int mapWidth;
    private int mapHeight;
    private FiveOpen fiveOpen;

    public MTCLTreeSearch(int simulateCount)
    {
        this.simulateCount = simulateCount;
    }

    public Node Select(Node node)
    {
        // 未完全展开
        if (!node.fullExpend)
        {
            var chNode = new Node()
            {
                averageValue = 0,
                chiList = new List<Node>(),
                moveData = node.notVisitList[0],
                notVisitList = new List<MoveData>(),
                parent = node,
                value = 0,
                visitCount = 0,
                winCount = 0,
                color = -node.color,
                depth = node.depth + 1,
                fullExpend = false
            };

            node.notVisitList.RemoveAt(0);
            node.fullExpend = node.notVisitList.Count == 0;

            node.chiList.Add(chNode);

            return chNode;
        }

        // 完全展开
        float maxUcb = float.MinValue;
        Node maxNode = null;
        foreach (var childNode in node.chiList)
        {
            float ucb = childNode.averageValue + 1.414f * Mathf.Sqrt(Mathf.Log(node.visitCount) / node.visitCount);
            if (ucb > maxUcb)
            {
                maxNode = childNode;
                maxUcb = ucb;
            }
        }

        return maxNode;
    }

    public void Expend(Node node)
    {
        for (int i = 0; i < mapWidth; i++)
        {
            for (int j = 0; j < mapHeight; j++)
            {
                if (fiveOpen.map[i, j] == 0)
                {
                    var moveData = new MoveData() { x = i, y = j };
                    node.notVisitList.Add(moveData);

                    // 检查是否游戏结束
                    fiveOpen.map[i, j] = -node.color;
                    var res = fiveOpen.CheckResult();
                    fiveOpen.map[i, j] = 0;
                    if (res != ChessResult.Incomplete)
                    {
                        node.notVisitList.Clear();
                        node.notVisitList.Add(moveData);
                        return;
                    }
                }
            }
        }
    }

    public int Simulate(FiveOpen fiveOpen)
    {
        List<MoveData> moveDatas = new List<MoveData>();
        for (int i = 0; i < mapWidth; i++)
        {
            for (int j = 0; j < mapHeight; j++)
            {
                if (fiveOpen.map[i, j] == 0)
                {
                    var md = new MoveData() { x = i, y = j };
                    fiveOpen.map[i, j] = fiveOpen.currentPlayer;
                    var gameRes = fiveOpen.CheckResult();
                    fiveOpen.map[i, j] = 0;
                    if (gameRes != ChessResult.Incomplete)
                    {
                        return (int)gameRes;
                    }

                    moveDatas.Add(md);
                }
            }
        }

        var moveData = moveDatas[Random.Range(0, moveDatas.Count)];
        fiveOpen.PlacePiece(moveData.x, moveData.y);
        fiveOpen.TurnPlayer();

        return Simulate(fiveOpen);
    }

    public void Back(Node node, int winRes)
    {
        node.visitCount++;
        if (winRes == node.color)
        {
            node.winCount++;
            node.value++;
        }
        else if (winRes != 0)
        {
            node.value--;
        }

        node.averageValue = node.value / node.visitCount;

        if (node.parent != null)
        {
            Back(node.parent, winRes);
        }
    }

    public void Search(Node node)
    {
        if (node.notVisitList.Count == 0 && !node.fullExpend)
        {
            Expend(node);
            if (node.notVisitList.Count == 0)
            {
                return;
            }

            Search(node);
        }
        else
        {
            var curNode = Select(node);
            fiveOpen.map[curNode.moveData.x, curNode.moveData.y] = curNode.color;
            var gameRes = fiveOpen.CheckResult();
            if (gameRes != ChessResult.Incomplete)
            {
                Back(curNode, (int)gameRes);
            }
            else if (curNode.visitCount == 0)
            {
                int winRes = Simulate(fiveOpen.Clone());
                Back(curNode, winRes);
            }
            else
            {
                Search(curNode);
            }

            fiveOpen.map[curNode.moveData.x, curNode.moveData.y] = 0;
        }
    }

    public MoveData Search(FiveOpen fiveOpen)
    {
        mapWidth = fiveOpen.map.GetLength(0);
        mapHeight = fiveOpen.map.GetLength(1);
        this.fiveOpen = fiveOpen;

        Node root = new Node()
        {
            chiList = new List<Node>(),
            moveData = null,
            parent = null,
            visitCount = 1,
            winCount = 0,
            value = 0,
            averageValue = 0,
            depth = 0,
            color = -fiveOpen.currentPlayer,
            notVisitList = new List<MoveData>()
        };

        for (int i = 0; i < simulateCount; i++)
        {
            Search(root);
        }

        return Select(root).moveData;
    }
}

public class Node
{
    public MoveData moveData;
    public Node parent;
    public List<Node> chiList;
    public List<MoveData> notVisitList;
    public float visitCount;
    public float winCount;
    public float value;
    public float averageValue;
    public int depth;
    public int color;
    public bool fullExpend;
}

public class MoveData
{
    public int x;
    public int y;
}